from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import Dict, Any

mcp = FastMCP("TensorFlow MCP Server")

"""
TensorFlow MCP Server

面向软件：tensorflow (Python 机器学习框架)

MCP tools 列表：
1. tensorflow_python_test
   - 用途：运行 TensorFlow Python 测试脚本
   - 参数：无
   - 返回结构：统一的 JSON 结果，包含命令执行状态、退出码、标准输出和错误输出
"""

@mcp.tool()
def tensorflow_python_test(ctx: Context) -> Dict[str, Any]:
    """
    运行 TensorFlow Python 测试脚本
    
    用途：执行 TensorFlow 的 Python 测试脚本，验证 TensorFlow 基本功能是否正常
    参数：无
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的完整命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 标准输出内容
        "stderr": str         # 错误输出内容
    }
    """
    try:
        # 执行测试命令
        cmd = "python3 common/test_python3-tensorflow.py | grep 'Assertion passed'"
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=300)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 300 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Exception occurred: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()